// Package ginlog
/**
* @Project : geektime-basic-go-study
* @File    : builder.go
* @IDE     : GoLand
* @Author  : Tvux
* @Date    : 2024/11/6 00:17
**/

package ginlog

import (
	"bytes"
	"go.uber.org/atomic"
	"io"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/spf13/viper"
	"golang.org/x/net/context"
)

// Builder 注意点：
// 1、小心日志内容过多。URL 可能很长，请求体，响应体都可能很大，你要考虑是不是完全输出到日志里面
// 2、考虑 1 的问题，以及用户可能换用不同的日志框架，所以要有足够的灵活性
// 3、考虑动态开关，结合监听配置文件，要小心并发安全
type Builder struct {
	allowReqBody  *atomic.Bool
	allowRespBody *atomic.Bool
	loggerFunc    func(ctx context.Context, al *AccessLog)
}

func NewBuilder(fn func(ctx context.Context, al *AccessLog)) *Builder {
	builder := &Builder{
		loggerFunc:    fn,
		allowReqBody:  atomic.NewBool(false),
		allowRespBody: atomic.NewBool(false),
	}
	builder.allowReqBody.Store(viper.GetBool("config.log.allowReqBody"))
	builder.allowRespBody.Store(viper.GetBool("config.log.allowRespBody"))
	return builder
}

func (Self *Builder) AllowReqBody(ok bool) *Builder {
	Self.allowReqBody.Store(ok)
	return Self
}

func (Self *Builder) AllowRespBody(ok bool) *Builder {
	Self.allowRespBody.Store(ok)
	return Self
}

func (Self *Builder) Build() gin.HandlerFunc {
	return func(ctx *gin.Context) {
		start := time.Now()
		url := ctx.Request.URL.String()
		al := &AccessLog{
			Method: ctx.Request.Method,
			URL:    url,
		}
		if Self.allowReqBody.Load() && ctx.Request.Body != nil {
			// Body 读出来之后就没有了，因为是一个流，所以还要放回去
			body, _ := ctx.GetRawData()
			ctx.Request.Body = io.NopCloser(bytes.NewReader(body))
			al.ReqBody = string(body) // 注意：这个会引起复制，会消耗cpu跟内存
		}

		if Self.allowRespBody.Load() {
			// 因为 gin 本身没有提供读取响应的方法，所以只能自己写一个 ResponseWriter 来截取响应
			ctx.Writer = respWriter{
				al:             al,
				ResponseWriter: ctx.Writer,
			}
		}

		defer func() {
			// 最后打印日志，防止程序panic不记录日志
			al.Duration = time.Since(start).String()
			Self.loggerFunc(ctx, al)
		}()

		// 执行业务逻辑
		ctx.Next()
	}
}

// respWriter 是一个自定义的 ResponseWriter，用于截取响应并记录响应体。
type respWriter struct {
	al *AccessLog
	gin.ResponseWriter
}

func (Self respWriter) WriteHeader(statusCode int) {
	Self.al.StatusCode = statusCode
	Self.ResponseWriter.WriteHeader(statusCode)
}

func (Self respWriter) Write(data []byte) (int, error) {
	Self.al.RespBody = string(data)
	return Self.ResponseWriter.Write(data)
}

func (Self respWriter) WriteString(data string) (int, error) {
	Self.al.RespBody = data
	return Self.ResponseWriter.WriteString(data)
}

// AccessLog 记录访问日志的结构体。
type AccessLog struct {
	// HTTP 请求的方法
	Method string
	// URL 整个请求URL
	URL        string
	Duration   string
	ReqBody    string
	RespBody   string
	StatusCode int
}
